#include "TwoDimPlot.h"

#include "TH2F.h"
#include "TCanvas.h"

void DrawTwoDimPlot(std::string name, YAML::Node config, TChain* chain)
{
    //Histogram
    std::pair<int, int> bins = std::pair<int, int>(config["Bins"][0].as<int>(50), config["Bins"][1].as<int>(50));
    std::pair<double, double> xRange = std::pair<double, double>(config["XRange"][0].as<double>(0), config["XRange"][1].as<double>(5));
    std::pair<double, double> yRange = std::pair<double, double>(config["YRange"][0].as<double>(0), config["YRange"][1].as<double>(5));
    TH2F* thTwoDim = new TH2F(name.c_str(), "", bins.first, xRange.first, xRange.second, bins.second, yRange.first, yRange.second);
    std::string xName = config["XName"].as<std::string>(name);
    std::string yName = config["YName"].as<std::string>(name);
    std::stringstream projectStream;
    projectStream << yName << ":" << xName;
    std::string cut = config["Cut"].as<std::string>("1");
    chain->Project(name.c_str(), projectStream.str().c_str(), cut.c_str());

    //Title
    std::stringstream xTitleStream, yTitleStream;
    xTitleStream << config["XTitle"].as<std::string>(xName);
    yTitleStream << config["YTitle"].as<std::string>(yName);
    thTwoDim->GetXaxis()->SetTitle(xTitleStream.str().c_str());
    thTwoDim->GetYaxis()->SetTitle(yTitleStream.str().c_str());
    thTwoDim->GetYaxis()->SetMaxDigits(3);

    //Draw
    TCanvas* can = new TCanvas("can", "", 800, 600);
    TVirtualPad* pad = can->cd();
    std::string drawOption = config["DrawOption"].as<std::string>("SCAT");
    if (drawOption.find("z") != std::string::npos || drawOption.find("Z") != std::string::npos)
    {
        pad->SetTopMargin(0.08);
        pad->SetLeftMargin(0.15);
        pad->SetRightMargin(0.15);
    }
    thTwoDim->Draw(drawOption.c_str());

    //Save
    std::stringstream plotStream;
    std::string plotType = config["PlotType"].as<std::string>("eps");
    plotStream << "ezDraw_TwoDim_" << name << "." << plotType;
    std::string plotName = config["PlotName"].as<std::string>(plotStream.str());
    can->SaveAs(plotName.c_str());

    delete thTwoDim;
    delete can;
}
